package net.aiion.weave.impl.injection.transformation;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

import net.aiion.weave.asm.Label;
import net.aiion.weave.asm.Opcodes;
import net.aiion.weave.asm.Type;
import net.aiion.weave.spi.bytecode.Bytecode;
import net.aiion.weave.spi.bytecode.BytecodeContext;
import net.aiion.weave.spi.bytecode.BytecodeException;
import net.aiion.weave.spi.bytecode.BytecodeGenerator;
import net.aiion.weave.spi.bytecode.events.AnnotationEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.AnnotationStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.AnnotationValueBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.BytecodeEvent;
import net.aiion.weave.spi.bytecode.events.ClassEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.ClassStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.MethodEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.MethodStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.util.BytecodeTransformation;
import net.aiion.weave.spi.injection.InjectionContext;
import net.aiion.weave.spi.injection.Injector;
import net.aiion.weave.spi.injection.annotations.Injects;

// TODO Add inheritance support (also respect annotations of implemented interfaces).
// TODO Add real support for multiple @Inject annotated methods. Currently, a call to InjectionContext.forwardInjection() skips all other @Injects annotated methods for the same class.
public class InjectorTransformation extends BytecodeTransformation
{
    private static final String INJECTION_CONTEXT_INTERNAL_NAME = Type.getInternalName(InjectionContext.class);

    private String className;

    private String methodName;
    private String methodDescriptor;

    private Type injectedAnnotation;

    private boolean inInjectsAnnotation;

    private Collection<BytecodeEvent> injectValueBytecodeEvents;
    private BytecodeGenerator injectValueBytecodeGenerator;

    private Collection<Collection<BytecodeEvent>> injectValueCodeBytecodeEvents;
    private Map<Type, Collection<BytecodeEvent>> injectValueByAnnotationsBytecodeEvents;
    private Map<Type, BytecodeGenerator> injectValueByAnnotationsBytecodeGenerators;

    private Label nextAnnotationLabel;
    private Map<Type, Label> nextTargetTypeLabelByAnnotationName;

    public InjectorTransformation() {
    }

    @Override
    public boolean testRetain(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.CLASS_START
            && !aBytecodeEvent.asClassStart().containsInterface(Injector.class);
    }

    @Override
    public boolean testProcess(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.ANNOTATION_START
            && aBytecodeEvent.asAnnotationStart().equalsType(Injects.class);
    }

    @Override
    public boolean testSkip(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.CLASS_START
            && (aBytecodeEvent.asClassStart().isInterface() || aBytecodeEvent.asClassStart().containsInterface(
                Injector.class));
    }

    @Override
    public void handleBytecodeEvent(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        switch (aBytecodeEvent.getType()) {
        case BytecodeEvent.CLASS_START:
            handleClassStartBytecode(aBytecodeContext, aBytecodeEvent.asClassStart());
            break;
        case BytecodeEvent.METHOD_START:
            handleMethodStartBytecode(aBytecodeContext, aBytecodeEvent.asMethodStart());
            break;
        case BytecodeEvent.ANNOTATION_START:
            handleAnnotationStartBytecode(aBytecodeContext, aBytecodeEvent.asAnnotationStart());
            break;
        case BytecodeEvent.ANNOTATION_VALUE:
            handleAnnotationValueBytecode(aBytecodeContext, aBytecodeEvent.asAnnotationValue());
            break;
        case BytecodeEvent.ANNOTATION_END:
            handleAnnotationEndBytecode(aBytecodeContext, aBytecodeEvent.asAnnotationEnd());
            break;
        case BytecodeEvent.METHOD_END:
            handleMethodEndBytecode(aBytecodeContext, aBytecodeEvent.asMethodEnd());
            break;
        case BytecodeEvent.CLASS_END:
            handleClassEndBytecode(aBytecodeContext, aBytecodeEvent.asClassEnd());
            break;
        default:
            break;
        }

        aBytecodeContext.forwardBytecodeEvent(aBytecodeEvent);
    }

    private void handleClassStartBytecode(final BytecodeContext aBytecodeContext,
        final ClassStartBytecodeEvent aClassStartBytecodeEvent)
    {
        className = aClassStartBytecodeEvent.getName();
        aClassStartBytecodeEvent.addInterface(Injector.class);

        prepareInjectValue();
    }

    private void handleMethodStartBytecode(final BytecodeContext aBytecodeContext,
        final MethodStartBytecodeEvent aMethodStartBytecodeEvent)
    {
        methodName = aMethodStartBytecodeEvent.getName();
        methodDescriptor = aMethodStartBytecodeEvent.getDescriptor();
    }

    private void handleAnnotationStartBytecode(final BytecodeContext aBytecodeContext,
        final AnnotationStartBytecodeEvent aAnnotationStartBytecodeEvent)
    {
        if (methodName != null && aAnnotationStartBytecodeEvent.equalsType(Injects.class)) {
            inInjectsAnnotation = true;

            testInjectorMethod();
        }
    }

    private void testInjectorMethod() {
        if (!methodDescriptor.startsWith("(L" + INJECTION_CONTEXT_INTERNAL_NAME + ";)")
            || methodDescriptor.endsWith(")V"))
        {
            throw new BytecodeException("A @Injects annotated method must accept exactly one argument of type "
                + InjectionContext.class.getName() + " and must not return void.");
        }
    }

    private void handleAnnotationValueBytecode(final BytecodeContext aBytecodeContext,
        final AnnotationValueBytecodeEvent aAnnotationValueBytecodeEvent)
    {
        if (inInjectsAnnotation && "value".equals(aAnnotationValueBytecodeEvent.getName())) {
            injectedAnnotation = (Type) aAnnotationValueBytecodeEvent.getValue();
        }
    }

    private void handleAnnotationEndBytecode(final BytecodeContext aBytecodeContext,
        final AnnotationEndBytecodeEvent aAnnotationEndBytecodeEvent)
    {
        inInjectsAnnotation = false;
    }

    private void handleMethodEndBytecode(final BytecodeContext aBytecodeContext,
        final MethodEndBytecodeEvent aMethodEndBytecodeEvent)
    {
        if (injectedAnnotation != null) {
            extendInjectValue();
        }

        methodName = null;
        injectedAnnotation = null;
    }

    private void handleClassEndBytecode(final BytecodeContext aBytecodeContext,
        final ClassEndBytecodeEvent aClassEndBytecodeEvent)
    {
        finishInjectValue();

        aBytecodeContext.forwardBytecodeEvents(injectValueBytecodeEvents);
    }

    private void prepareInjectValue() {
        injectValueBytecodeEvents = new ArrayList<>();
        injectValueBytecodeGenerator = Bytecode.createGenerator(injectValueBytecodeEvents);

        injectValueBytecodeGenerator.startMethod(Opcodes.ACC_PUBLIC, "injectValue", Object.class,
            InjectionContext.class);
        injectValueBytecodeGenerator.startCode();

        // <<2>> = <<1>>.getTargetType()
        injectValueBytecodeGenerator.loadReference(1);
        injectValueBytecodeGenerator.invokeInterface(InjectionContext.class, "getTargetType", Class.class);
        injectValueBytecodeGenerator.storeReference(2);

        injectValueCodeBytecodeEvents = new ArrayList<>();
        injectValueByAnnotationsBytecodeEvents = new HashMap<>();
        injectValueByAnnotationsBytecodeGenerators = new HashMap<>();
        nextTargetTypeLabelByAnnotationName = new HashMap<>();
    }

    private void extendInjectValue() {
        final boolean isFirstAnnotation = injectValueByAnnotationsBytecodeEvents.isEmpty();

        if (!injectValueByAnnotationsBytecodeGenerators.containsKey(injectedAnnotation)) {
            final Collection<BytecodeEvent> bytecodeEvents = new ArrayList<>();
            final BytecodeGenerator bytecodeGenerator = Bytecode.createGenerator(bytecodeEvents);

            injectValueCodeBytecodeEvents.add(bytecodeEvents);
            injectValueByAnnotationsBytecodeEvents.put(injectedAnnotation, bytecodeEvents);
            injectValueByAnnotationsBytecodeGenerators.put(injectedAnnotation, bytecodeGenerator);
        }

        final Collection<BytecodeEvent> bytecodeEvents = injectValueByAnnotationsBytecodeEvents.get(injectedAnnotation);
        final BytecodeGenerator bytecodeGenerator = injectValueByAnnotationsBytecodeGenerators.get(injectedAnnotation);

        if (bytecodeEvents.isEmpty()) {
            if (!isFirstAnnotation) {
                generateJumpOrReturnReference3(bytecodeGenerator, nextAnnotationLabel);

                bytecodeGenerator.markLabel(nextAnnotationLabel);
                bytecodeGenerator.markFrameChop(1, null, 0, null);
            }

            nextAnnotationLabel = new Label();
            nextTargetTypeLabelByAnnotationName.put(injectedAnnotation, new Label());

            // <<1>>.hasTargetAnnotation(<<targetAnnotation>>) IFEQ <<nextAnnotationLabel>>
            bytecodeGenerator.loadReference(1);
            bytecodeGenerator.pushConstant(injectedAnnotation);
            bytecodeGenerator
                .invokeInterface(InjectionContext.class, "hasTargetAnnotation", boolean.class, Class.class);
            bytecodeGenerator.jumpEquals(nextAnnotationLabel);

            generateInvokeDelegateAndTestTargetType(bytecodeGenerator);
        } else {
            final Label nextTargetTypeLabel = nextTargetTypeLabelByAnnotationName.get(injectedAnnotation);

            generateJumpOrReturnReference3(bytecodeGenerator, nextTargetTypeLabel);

            bytecodeGenerator.markLabel(nextTargetTypeLabel);
            bytecodeGenerator.markFrameAppend(2, new Object[] {
                "java/lang/Class", "java/lang/Object"
            }, 0, null);

            generateInvokeDelegateAndTestTargetType(bytecodeGenerator);
        }
    }

    private void finishInjectValue() {
        for (Collection<BytecodeEvent> bytecodeEvents : injectValueCodeBytecodeEvents) {
            injectValueBytecodeEvents.addAll(bytecodeEvents);
        }

        generateJumpOrReturnReference3(injectValueBytecodeGenerator, nextAnnotationLabel);

        injectValueBytecodeGenerator.markLabel(nextAnnotationLabel);
        injectValueBytecodeGenerator.markFrameSame();

        // return <<1>>.forwardInjection()
        injectValueBytecodeGenerator.loadReference(1);
        injectValueBytecodeGenerator.invokeInterface(InjectionContext.class, "forwardInjection", Object.class);
        injectValueBytecodeGenerator.returnReference();

        injectValueBytecodeGenerator.endCode(2, 4);
        injectValueBytecodeGenerator.endMethod();
    }

    private void generateInvokeDelegateAndTestTargetType(final BytecodeGenerator aBytecodeGenerator) {
        // <<3>> = this.createObject2(<<1>>)
        aBytecodeGenerator.loadReference(0);
        aBytecodeGenerator.loadReference(1);
        aBytecodeGenerator.invokeVirtual(className, methodName, methodDescriptor);
        aBytecodeGenerator.storeReference(3);

        // <<2>>.isInstance(<<3>>)
        aBytecodeGenerator.loadReference(2);
        aBytecodeGenerator.loadReference(3);
        aBytecodeGenerator.invokeVirtual(Class.class, "isInstance", boolean.class, Object.class);
    }

    private void generateJumpOrReturnReference3(final BytecodeGenerator aBytecodeGenerator, final Label aLabel) {
        // IFEQ <<aLabel>>
        // return <<3>>
        aBytecodeGenerator.jumpEquals(aLabel);
        aBytecodeGenerator.loadReference(3);
        aBytecodeGenerator.returnReference();
    }
}
